#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <geometry_msgs/PoseStamped.h>
#include <tf2/LinearMath/Transform.h>
#include <tf2_ros/transform_broadcaster.h>
#include <tf2_ros/transform_listener.h>
#include <tf2/LinearMath/Quaternion.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/registration/ndt.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/io/pcd_io.h>
#include <string>
#include <fstream>
#include <vector>
#include <cmath>
#include <ctime>
#include <iomanip>
#include <sstream>
#include <sys/stat.h>
#include <sys/types.h>

class NDTMapper
{
private:
    // ROS节点句柄
    ros::NodeHandle nh_;
    ros::NodeHandle private_nh_;

    // 订阅者和发布者
    ros::Subscriber points_sub_;
    ros::Publisher map_pub_;
    ros::Publisher current_pose_pub_;
    ros::Publisher current_scan_pub_;

    // TF广播器
    tf2_ros::TransformBroadcaster tf_broadcaster_;
    tf2_ros::Buffer tf_buffer_;
    tf2_ros::TransformListener tf_listener_;

    // 点云和位姿
    pcl::PointCloud<pcl::PointXYZ>::Ptr map_cloud_;
    pcl::PointCloud<pcl::PointXYZ>::Ptr current_cloud_;
    tf2::Transform current_pose_;
    bool is_first_scan_;

    // NDT参数
    double voxel_leaf_size_;
    double ndt_resolution_;
    double ndt_step_size_;
    double ndt_epsilon_;
    int max_iterations_;

    // 坐标系参数
    std::string map_frame_id_;
    std::string base_frame_id_;
    std::string lidar_frame_id_;

    // 地图保存参数
    std::string map_directory_;
    std::string timestamp_directory_;
    std::string default_map_name_;
    bool auto_save_on_shutdown_;
    bool auto_save_with_timestamp_;
    double auto_save_interval_;
    ros::Time last_save_time_;
    ros::Timer save_timer_;
    
    // 点云添加控制参数
    double cloud_add_interval_;  // 点云添加到地图的时间间隔（秒）
    ros::Time last_cloud_add_time_;  // 上次添加点云到地图的时间

public:
    NDTMapper() : private_nh_("~"), tf_listener_(tf_buffer_), is_first_scan_(true)
    {
        // 初始化点云
        map_cloud_.reset(new pcl::PointCloud<pcl::PointXYZ>());
        current_cloud_.reset(new pcl::PointCloud<pcl::PointXYZ>());

        // 获取参数
        private_nh_.param<double>("voxel_leaf_size", voxel_leaf_size_, 0.2);
        private_nh_.param<double>("ndt_resolution", ndt_resolution_, 1.0);
        private_nh_.param<double>("ndt_step_size", ndt_step_size_, 0.1);
        private_nh_.param<double>("ndt_epsilon", ndt_epsilon_, 0.01);
        private_nh_.param<int>("max_iterations", max_iterations_, 30);

        private_nh_.param<std::string>("map_frame_id", map_frame_id_, "map");
        private_nh_.param<std::string>("base_frame_id", base_frame_id_, "base_link");
        private_nh_.param<std::string>("lidar_frame_id", lidar_frame_id_, "velodyne");

        // 获取地图保存参数 - 修改为从私有命名空间加载
        private_nh_.param<std::string>("map_directory", map_directory_, "/root/catkin_ws/src/auto_navigation/maps");
        private_nh_.param<std::string>("default_map_name", default_map_name_, "map.pcd");
        private_nh_.param<bool>("auto_save_on_shutdown", auto_save_on_shutdown_, true);
        private_nh_.param<bool>("auto_save_with_timestamp", auto_save_with_timestamp_, true);
        private_nh_.param<double>("auto_save_interval", auto_save_interval_, 30.0); // 默认30秒保存一次
        
        // 获取点云添加间隔参数
        private_nh_.param<double>("cloud_add_interval", cloud_add_interval_, 30.0); // 默认30秒添加一次点云
        last_cloud_add_time_ = ros::Time::now();

        // 添加更多调试信息
        ROS_INFO("NDT Mapper parameters:");
        ROS_INFO("  voxel_leaf_size: %.2f", voxel_leaf_size_);
        ROS_INFO("  ndt_resolution: %.2f", ndt_resolution_);
        ROS_INFO("  ndt_step_size: %.2f", ndt_step_size_);
        ROS_INFO("  ndt_epsilon: %.3f", ndt_epsilon_);
        ROS_INFO("  max_iterations: %d", max_iterations_);
        ROS_INFO("  map_directory: %s", map_directory_.c_str());
        ROS_INFO("  auto_save_interval: %.1f", auto_save_interval_);
        ROS_INFO("  auto_save_with_timestamp: %s", auto_save_with_timestamp_ ? "true" : "false");
        ROS_INFO("  cloud_add_interval: %.1f", cloud_add_interval_);

        // 创建基于时间戳的目录
        createTimestampDirectory();

        // 初始化位姿
        current_pose_.setIdentity();

        // 订阅点云话题
        points_sub_ = nh_.subscribe("velodyne_points", 10, &NDTMapper::pointsCallback, this);

        // 发布地图和位姿
        map_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("ndt/map", 1, true);
        current_pose_pub_ = nh_.advertise<geometry_msgs::PoseStamped>("ndt/current_pose", 1, true);
        current_scan_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("ndt/current_scan", 1, true);

        // 如果启用了定时保存，设置定时器
        if (auto_save_interval_ > 0.0) {
            save_timer_ = nh_.createTimer(ros::Duration(auto_save_interval_), 
                                         &NDTMapper::saveTimerCallback, this);
            ROS_INFO("Auto save timer enabled, interval: %.1f seconds", auto_save_interval_);
        }

        ROS_INFO("NDT Mapper initialized");
        ROS_INFO("Auto save with timestamp: %s", auto_save_with_timestamp_ ? "enabled" : "disabled");
        ROS_INFO("Maps will be saved to: %s", timestamp_directory_.c_str());
    }

    ~NDTMapper()
    {
        if (auto_save_on_shutdown_)
        {
            if (auto_save_with_timestamp_) {
                saveMapWithTimestamp();
            } else {
                saveMap(timestamp_directory_ + "/" + default_map_name_);
            }
        }
    }

    void createTimestampDirectory()
    {
        // 获取当前时间
        auto now = std::chrono::system_clock::now();
        auto in_time_t = std::chrono::system_clock::to_time_t(now);
        
        std::stringstream ss;
        ss << map_directory_ << "/";
        ss << std::put_time(std::localtime(&in_time_t), "%Y%m%d_%H%M%S");
        
        timestamp_directory_ = ss.str();
        
        // 创建目录
        std::string command = "mkdir -p " + timestamp_directory_;
        int result = system(command.c_str());
        
        if (result == 0) {
            ROS_INFO("Created timestamp directory: %s", timestamp_directory_.c_str());
            
            // 创建符号链接到最新的地图目录
            std::string latest_dir_link = map_directory_ + "/latest_mapping";
            std::string cmd = "ln -sfn \"" + timestamp_directory_ + "\" \"" + latest_dir_link + "\"";
            system(cmd.c_str());
            ROS_INFO("Created symbolic link: %s -> %s", latest_dir_link.c_str(), timestamp_directory_.c_str());
        } else {
            ROS_ERROR("Failed to create timestamp directory: %s", timestamp_directory_.c_str());
            // 如果创建失败，使用默认目录
            timestamp_directory_ = map_directory_;
        }
    }

    void saveTimerCallback(const ros::TimerEvent& event)
    {
        if (map_cloud_->empty()) {
            ROS_WARN("Map is empty, skipping timed save.");
            return;
        }
        
        // 添加调试信息
        ROS_INFO("Auto save timer triggered at %.2f, saving map...", ros::Time::now().toSec());
        
        // 保存地图
        saveMapWithTimestamp();
        
        // 记录最后保存时间
        last_save_time_ = ros::Time::now();
    }

    void pointsCallback(const sensor_msgs::PointCloud2ConstPtr& msg)
    {
        // 将ROS消息转换为PCL点云
        pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>());
        pcl::fromROSMsg(*msg, *cloud);

        // 体素滤波降采样
        pcl::VoxelGrid<pcl::PointXYZ> voxel_filter;
        voxel_filter.setInputCloud(cloud);
        voxel_filter.setLeafSize(voxel_leaf_size_, voxel_leaf_size_, voxel_leaf_size_);
        voxel_filter.filter(*current_cloud_);

        // 如果是第一帧，直接添加到地图
        if (is_first_scan_)
        {
            *map_cloud_ += *current_cloud_;
            is_first_scan_ = false;
            last_cloud_add_time_ = ros::Time::now();
            publishMapAndPose();
            return;
        }

        // 使用NDT算法进行配准
        pcl::NormalDistributionsTransform<pcl::PointXYZ, pcl::PointXYZ> ndt;
        ndt.setTransformationEpsilon(ndt_epsilon_);
        ndt.setStepSize(ndt_step_size_);
        ndt.setResolution(ndt_resolution_);
        ndt.setMaximumIterations(max_iterations_);

        ndt.setInputSource(current_cloud_);
        ndt.setInputTarget(map_cloud_);

        // 使用当前位姿作为初始猜测
        Eigen::Matrix4f init_guess = Eigen::Matrix4f::Identity();
        tf2::Matrix3x3 rot_mat(current_pose_.getRotation());
        Eigen::Matrix3f eigen_rot;
        eigen_rot << rot_mat[0][0], rot_mat[0][1], rot_mat[0][2],
                     rot_mat[1][0], rot_mat[1][1], rot_mat[1][2],
                     rot_mat[2][0], rot_mat[2][1], rot_mat[2][2];
        init_guess.block<3, 3>(0, 0) = eigen_rot;
        init_guess(0, 3) = current_pose_.getOrigin().x();
        init_guess(1, 3) = current_pose_.getOrigin().y();
        init_guess(2, 3) = current_pose_.getOrigin().z();

        // 执行配准
        pcl::PointCloud<pcl::PointXYZ>::Ptr aligned_cloud(new pcl::PointCloud<pcl::PointXYZ>());
        ndt.align(*aligned_cloud, init_guess);

        // 输出NDT得分
        ROS_INFO("NDT score: %f", ndt.getFitnessScore());

        // 如果配准成功，更新位姿和地图
        if (ndt.hasConverged())
        {
            // 获取变换矩阵
            Eigen::Matrix4f transform = ndt.getFinalTransformation();

            // 更新当前位姿
            tf2::Matrix3x3 tf_rot;
            tf_rot.setValue(
                transform(0, 0), transform(0, 1), transform(0, 2),
                transform(1, 0), transform(1, 1), transform(1, 2),
                transform(2, 0), transform(2, 1), transform(2, 2)
            );
            tf2::Quaternion tf_quat;
            tf_rot.getRotation(tf_quat);
            current_pose_.setOrigin(tf2::Vector3(transform(0, 3), transform(1, 3), transform(2, 3)));
            current_pose_.setRotation(tf_quat);

            // 检查是否达到点云添加间隔
            ros::Time current_time = ros::Time::now();
            double time_since_last_add = (current_time - last_cloud_add_time_).toSec();
            
            // 只有当时间间隔超过设定值时，才添加点云到地图
            if (time_since_last_add >= cloud_add_interval_)
            {
                ROS_INFO("Adding cloud to map after %.1f seconds", time_since_last_add);
                
                // 将当前点云添加到地图
                *map_cloud_ += *aligned_cloud;
                
                // 更新最后添加时间
                last_cloud_add_time_ = current_time;
                
                // 检查地图大小，如果过大则降采样
                if (map_cloud_->size() > 5000000) { // 500万点限制
                    ROS_WARN("Map size exceeded limit (%ld points), downsampling...", map_cloud_->size());
                    pcl::PointCloud<pcl::PointXYZ>::Ptr temp_cloud(new pcl::PointCloud<pcl::PointXYZ>());
                    pcl::VoxelGrid<pcl::PointXYZ> voxel_grid;
                    voxel_grid.setInputCloud(map_cloud_);
                    voxel_grid.setLeafSize(voxel_leaf_size_ * 2, voxel_leaf_size_ * 2, voxel_leaf_size_ * 2);
                    voxel_grid.filter(*temp_cloud);
                    map_cloud_ = temp_cloud;
                    ROS_INFO("Map downsampled to %ld points", map_cloud_->size());
                }
            }
            else
            {
                // 如果时间间隔不够，只更新位姿，不添加点云到地图
                ROS_DEBUG("Skipping cloud addition, %.1f seconds since last add (interval: %.1f)", 
                         time_since_last_add, cloud_add_interval_);
            }

            // 发布地图和位姿
            publishMapAndPose();
        }
    }

    void publishMapAndPose()
    {
        // 发布地图点云
        sensor_msgs::PointCloud2 map_msg;
        pcl::toROSMsg(*map_cloud_, map_msg);
        map_msg.header.stamp = ros::Time::now();
        map_msg.header.frame_id = map_frame_id_;
        map_pub_.publish(map_msg);

        // 发布当前点云
        sensor_msgs::PointCloud2 current_msg;
        pcl::toROSMsg(*current_cloud_, current_msg);
        current_msg.header.stamp = ros::Time::now();
        current_msg.header.frame_id = map_frame_id_;
        current_scan_pub_.publish(current_msg);

        // 发布当前位姿
        geometry_msgs::PoseStamped pose_msg;
        pose_msg.header.stamp = ros::Time::now();
        pose_msg.header.frame_id = map_frame_id_;
        pose_msg.pose.position.x = current_pose_.getOrigin().x();
        pose_msg.pose.position.y = current_pose_.getOrigin().y();
        pose_msg.pose.position.z = current_pose_.getOrigin().z();
        pose_msg.pose.orientation.x = current_pose_.getRotation().x();
        pose_msg.pose.orientation.y = current_pose_.getRotation().y();
        pose_msg.pose.orientation.z = current_pose_.getRotation().z();
        pose_msg.pose.orientation.w = current_pose_.getRotation().w();
        current_pose_pub_.publish(pose_msg);

        // 广播TF变换
        geometry_msgs::TransformStamped transform_stamped;
        transform_stamped.header.stamp = ros::Time::now();
        transform_stamped.header.frame_id = map_frame_id_;
        transform_stamped.child_frame_id = base_frame_id_;
        transform_stamped.transform.translation.x = current_pose_.getOrigin().x();
        transform_stamped.transform.translation.y = current_pose_.getOrigin().y();
        transform_stamped.transform.translation.z = current_pose_.getOrigin().z();
        transform_stamped.transform.rotation.x = current_pose_.getRotation().x();
        transform_stamped.transform.rotation.y = current_pose_.getRotation().y();
        transform_stamped.transform.rotation.z = current_pose_.getRotation().z();
        transform_stamped.transform.rotation.w = current_pose_.getRotation().w();
        tf_broadcaster_.sendTransform(transform_stamped);
    }

    std::string generateTimestampFilename()
    {
        // 获取当前时间
        auto now = std::chrono::system_clock::now();
        auto in_time_t = std::chrono::system_clock::to_time_t(now);
        
        std::stringstream ss;
        ss << timestamp_directory_ << "/kitti_map_";
        ss << std::put_time(std::localtime(&in_time_t), "%Y%m%d_%H%M%S");
        ss << ".pcd";
        
        return ss.str();
    }

    void saveMapWithTimestamp()
    {
        if (map_cloud_->empty())
        {
            ROS_WARN("Map is empty, not saving.");
            return;
        }

        std::string filename = generateTimestampFilename();
        saveMap(filename);
        
        // 创建符号链接到最新的地图
        std::string latest_link = timestamp_directory_ + "/kitti_latest_map.pcd";
        std::string cmd = "ln -sf \"" + filename + "\" \"" + latest_link + "\"";
        system(cmd.c_str());
        ROS_INFO("Created symbolic link: %s -> %s", latest_link.c_str(), filename.c_str());
        
        // 创建全局符号链接
        std::string global_latest_link = map_directory_ + "/kitti_latest_map.pcd";
        cmd = "ln -sf \"" + filename + "\" \"" + global_latest_link + "\"";
        system(cmd.c_str());
        ROS_INFO("Created global symbolic link: %s -> %s", global_latest_link.c_str(), filename.c_str());
    }

    void saveMap(const std::string& filename)
    {
        if (map_cloud_->empty())
        {
            ROS_WARN("Map is empty, not saving.");
            return;
        }

        // 确保目录存在
        std::string dir = filename.substr(0, filename.find_last_of('/'));
        std::string command = "mkdir -p " + dir;
        system(command.c_str());

        // 保存地图
        if (pcl::io::savePCDFile(filename, *map_cloud_) == 0)
        {
            ROS_INFO("Map saved to %s", filename.c_str());
        }
        else
        {
            ROS_ERROR("Failed to save map to %s", filename.c_str());
        }
    }
};

int main(int argc, char** argv)
{
    ros::init(argc, argv, "ndt_mapper_node");
    NDTMapper mapper;
    ros::spin();
    return 0;
} 